package com.henry.matrix;

import java.util.Arrays;

public class Matrix {
    int[] array;
    int rows;
    int cols;

    // public static String regexToMatchDelim = "\\s*,\\s*";

    public Matrix() {
        this.rows = 1;
        this.cols = 1;
        this.array = new int[1];
    }

    public Matrix(int nRows, int nCols) {
        this.array = new int[nRows * nCols];
        this.rows = nRows;
        this.cols = nCols;
        for (int i = 0; i < array.length; i++) { this.array[i] = 0; }
    }

    public Matrix(int nRows, int nCols, int[] inArray) {
        this.array = new int[nRows * nCols];
        this.rows = nRows;
        this.cols = nCols;
        for (int i = 0; i < array.length; i++) {
            if (i < inArray.length) this.array[i] = inArray[i];
            else this.array[i] = 0; // inArray has not enough elements, filling the rest with 0
        }
    }

    public Matrix (int nRows, int nCols, String inArray, String regex) {
        String str = inArray.trim();
        this.array = Arrays.stream(str.split(regex))
                .mapToInt(Integer::parseInt)
                .toArray();
        this.rows = nRows;
        this.cols = nCols;
    }

    public Matrix(Matrix other) {
        this.rows = other.rows;
        this.cols = other.cols;
        this.array = other.array.clone();
    }

    public Matrix add(Matrix other) {
        if (this.rows != other.rows || this.cols != other.rows) {
            throw new Error("Matrix shape does not match!");
        }
        int[] newArray = new int[array.length];
        for (int i = 0; i < newArray.length; ++i) {
            newArray[i] = this.array[i] + other.array[i];
        }
        return new Matrix(rows, cols, newArray);
    }

    public Matrix subtract(Matrix other) {
        if (this.rows != other.rows || this.cols != other.rows) {
            throw new Error("Matrix shape does not match!");
        }
        int[] newArray = new int[array.length];
        for (int i = 0; i < newArray.length; ++i) {
            newArray[i] = this.array[i] - other.array[i];
        }
        return new Matrix(rows, cols, newArray);
    }

    // A.multiply(B) means $A \cdot B$
    public Matrix multiply(Matrix other) {
        if (this.cols != other.rows) {
            throw new Error("Cannot multiply two matrices! A.cols != B.rows.");
        }
        int newRows = this.rows;
        int newCols = other.cols;
        int[] newArray = new int[newRows * newCols];
        for (int i = 0; i < newArray.length; i++) {
            int rowIdxOfA = i / newCols;
            int colIdxOfB = i % newCols;
            // M.array[i] = \Sigma A[rowIdxOfA][k] * B[k][colIdxOfB]
            for (int k = 0; k < this.cols; k++) {
                newArray[i] += this.array[rowIdxOfA * this.cols + k] * other.array[k * other.cols + colIdxOfB];
            }
        }
        return new Matrix(newRows, newCols, newArray);
    }

    public Matrix transpose() {
        int[] newArray = new int[array.length];
        for (int i = 0; i < array.length; i++) {
            int dataInOrgMatRow = i % rows;
            int dataInOrgMatCol = i / rows;
            newArray[i] = this.array[dataInOrgMatRow * cols + dataInOrgMatCol];
        }
        return new Matrix(cols, rows, newArray);
    }

    @Override
    public String toString() {
        StringBuilder out = new StringBuilder();
        for (int i = 0; i < rows * cols; ++i) {
            out.append(" ").append(array[i]);
            if (i % cols == cols - 1) out.append(";\n");
            else out.append(",");
        }
        out.setCharAt(0, '[');
        out.setCharAt(out.length() - 1, ']');
        return out.toString();
    }

    public int valueAt(int row, int col) { return array[row * rows + col]; }

    public void setValueAt(int row, int col, int newVal) { array[row * rows + col] = newVal; }

    public int[] data() { return this.array; }

    public int getRows() { return this.rows; }

    public int getCols() { return this.cols; }
}
